// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/parse/type_resolver.h>
#include <pollux/core/i_typed_expr.h>
#include <pollux/expression/function_call_to_special_form.h>
#include <pollux/expression/signature_binder.h>
#include <pollux/functions/function_registry.h>
#include <pollux/parse/expressions.h>
#include <pollux/type/type.h>

namespace kumo::pollux::parse {
    namespace {
        std::string toString(
            const std::string &functionName,
            const std::vector<TypePtr> &argTypes) {
            std::ostringstream signature;
            signature << functionName << "(";
            for (auto i = 0; i < argTypes.size(); i++) {
                if (i > 0) {
                    signature << ", ";
                }
                signature << argTypes[i]->toString();
            }
            signature << ")";
            return signature.str();
        }

        std::string toString(
            const std::vector<const exec::FunctionSignature *> &signatures) {
            std::stringstream out;
            for (auto i = 0; i < signatures.size(); ++i) {
                if (i > 0) {
                    out << ", ";
                }
                out << signatures[i]->toString();
            }
            return out.str();
        }

        TypePtr resolveType(
            const std::vector<std::shared_ptr<const core::ITypedExpr> > &inputs,
            const std::shared_ptr<const core::CallExpr> &expr,
            bool nullOnFailure) {
            std::vector<TypePtr> inputTypes;
            inputTypes.reserve(inputs.size());
            for (auto &input: inputs) {
                inputTypes.emplace_back(input->type());
            }

            if (auto resolvedType = exec::resolveTypeForSpecialForm(
                expr->getFunctionName(), inputTypes)) {
                return resolvedType;
            }

            return resolveScalarFunctionType(
                expr->getFunctionName(), inputTypes, nullOnFailure);
        }
    } // namespace

    void registerTypeResolver() {
        core::Expressions::setTypeResolverHook(&resolveType);
    }

    TypePtr resolveScalarFunctionType(
        const std::string &name,
        const std::vector<TypePtr> &argTypes,
        bool nullOnFailure) {
        auto returnType = resolve_function(name, argTypes);
        if (returnType) {
            return returnType;
        }

        if (nullOnFailure) {
            return nullptr;
        }

        auto allSignatures = get_function_signatures();
        auto it = allSignatures.find(name);
        if (it == allSignatures.end()) {
            POLLUX_USER_FAIL("Scalar function doesn't exist: {}.", name);
        } else {
            const auto &functionSignatures = it->second;
            POLLUX_USER_FAIL(
                "Scalar function signature is not supported: {}. Supported signatures: {}.",
                toString(name, argTypes),
                toString(functionSignatures));
        }
    }
} // namespace kumo::pollux::parse
